use super::{BNScalar, HyperKZGPublicSetupOwned};
use crate::{
    base::{
        proof::{Keccak256Transcript, Transcript},
        slice_ops,
    },
    proof_primitive::hyperkzg::convert_g1_affine_from_halo2_to_ark,
};
use nova_snark::{
    errors::NovaError,
    provider::{bn256_grumpkin::bn256::Scalar as NovaScalar, hyperkzg::CommitmentKey},
    traits::{Engine, TranscriptEngineTrait, TranscriptReprTrait},
};
use serde::{Deserialize, Serialize};

/// The `HyperKZG` engine that implements nova's `Engine` trait.
#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct HyperKZGEngine;

impl Engine for HyperKZGEngine {
    type Base = nova_snark::provider::bn256_grumpkin::bn256::Base;
    type Scalar = NovaScalar;
    type GE = nova_snark::provider::bn256_grumpkin::bn256::Point;
    type RO = nova_snark::provider::poseidon::PoseidonRO<Self::Base>;
    type ROCircuit = nova_snark::provider::poseidon::PoseidonROCircuit<Self::Base>;
    type RO2 = nova_snark::provider::poseidon::PoseidonRO<Self::Scalar>;
    type RO2Circuit = nova_snark::provider::poseidon::PoseidonROCircuit<Self::Scalar>;
    type TE = Keccak256Transcript;
    type CE = nova_snark::provider::hyperkzg::CommitmentEngine<Self>;
}

impl TranscriptEngineTrait<HyperKZGEngine> for Keccak256Transcript {
    fn new(_label: &'static [u8]) -> Self {
        Transcript::new()
    }

    fn squeeze(&mut self, _label: &'static [u8]) -> Result<NovaScalar, NovaError> {
        let res = Transcript::scalar_challenge_as_be::<BNScalar>(self).into();
        Transcript::challenge_as_le(self);
        Ok(res)
    }

    fn absorb<T: TranscriptReprTrait<<HyperKZGEngine as Engine>::GE>>(
        &mut self,
        _label: &'static [u8],
        o: &T,
    ) {
        Transcript::extend_as_le_from_refs(self, &o.to_transcript_bytes());
    }

    fn dom_sep(&mut self, _bytes: &'static [u8]) {}
}

/// Utility converting a nova `CommitmentKey` to a [`HyperKZGPublicSetupOwned`].
pub fn nova_commitment_key_to_hyperkzg_public_setup(
    setup: &CommitmentKey<HyperKZGEngine>,
) -> HyperKZGPublicSetupOwned {
    slice_ops::slice_cast_with(setup.ck(), convert_g1_affine_from_halo2_to_ark)
}
